import math
import time
import numpy as np
from thinkgear.serial import ThinkGearSerial
from thinkgear import *
import serial.tools.list_ports
import sys
import pyqtgraph as pg
from pyqtgraph.Qt import QtWidgets
import threading

class BMD101App:
    def __init__(self):
        self.app = pg.mkQApp("Plotting Example")
        self.win = QtWidgets.QMainWindow()
        self.win.setWindowTitle('BMD101')
         # Central widget
        central_widget = QtWidgets.QWidget()
        self.win.setCentralWidget(central_widget)

        # Layout
        layout = QtWidgets.QVBoxLayout()
        central_widget.setLayout(layout)

        # First row: Connect button and Status label
        toolbar_layout = QtWidgets.QHBoxLayout()

        self.connect_btn = QtWidgets.QPushButton("Connect")
        self.connect_btn.setMaximumWidth(100)
        self.connect_btn.clicked.connect(self.toggle_connection)
        toolbar_layout.addWidget(self.connect_btn)

        self.status_label = QtWidgets.QLabel("Not Connected")
        toolbar_layout.addWidget(self.status_label)

        self.frequency_label = QtWidgets.QLabel("Frequency: 0.0 Hz")
        toolbar_layout.addWidget(self.frequency_label)

        self.signal_quality_label = QtWidgets.QLabel("Signal Quality: 0")
        toolbar_layout.addWidget(self.signal_quality_label)

        layout.addLayout(toolbar_layout)
        # Second row: Plot
        self.signal_plot = pg.PlotWidget(title="Raw Signal")
        self.data_curve = self.signal_plot.plot(pen='w')
        self.local_voltage_curve =self.signal_plot.plot(pen='r')
        self.moving_avg_curve = self.signal_plot.plot(pen='y')
        self.local_max_curve = self.signal_plot.plot(symbolBrush='g')
        self.local_min_curve = self.signal_plot.plot(symbolBrush='b')
        layout.addWidget(self.signal_plot)


        self.t = None
        self.is_listening = False
        self.gain = 128
        self.conversion_factor = 0.0183
        # Data storage
        self.timestamps = []
        self.data = []
        self.signal_qualitys = []
        self.local_voltages = []
        self.moving_averages = []
        self.local_max = []
        self.local_min = []
        self.poor_signal = 0
        self.last_update_time = 0

    def run(self):
        self.win.showMaximized()
        self.app.exec()

    def toggle_connection(self):
        if self.is_listening:
            self.is_listening = False
        else:
            ports = serial.tools.list_ports.comports()
            available_ports = [port.device for port in ports]
            if len(available_ports) == 0:
                self.t = None
            else:
                self.t = ThinkGearSerial()
                self.status_label.setText(f"Connecting to {available_ports[0]}...")
                # Connect to the device (e.g., COM1 at 57600 baud)
                self.t.connect((available_ports[0], 115200))
                self.status_label.setText(f"{available_ports[0]} connected with baud rate {self.t.device.baudrate}")
                threading.Thread(target=self.listen, daemon=True).start()
                self.connect_btn.setText("Disconnect")

    def listen(self):
        if self.is_listening:
            return
        
        self.is_listening = True

        # Continuously read and process data
        while self.is_listening:
            ret = self.t.read()
            timestamp = time.time()            
            for data_point in ret:
                if isinstance(data_point, PoorSignalDataPoint):
                    self.poor_signal = data_point.value
                    self.signal_quality_label.setText(f"Signal Quality: {self.poor_signal}")
                elif isinstance(data_point, AttentionDataPoint):
                    print("Attention:", data_point.value)
                elif isinstance(data_point, MeditationDataPoint):
                    print("Meditation:", data_point.value)
                elif isinstance(data_point, RawDataPoint):
                    if data_point.code == 0x80:                        
                        self.update_data(timestamp, data_point.value)
                    else:
                        print(f"Data point: code: {data_point.code}, size: {len(data_point.data)}, data: {data_point.data}")

                    
        self.t.disconnect()
        self.t = None
        self.connect_btn.setText("Connect")
        self.status_label.setText("Not Connected")
    def update_data(self, timestamp, value):        
        self.timestamps.append(timestamp)     
        # value = np.clip(value, -10000, 10000)
        value = value / self.gain * self.conversion_factor
        self.data.append(value)
        self.signal_qualitys.append(self.poor_signal)
        self.update_local_attrs(window=1)
        if timestamp - self.last_update_time > 0.1:
            self.last_update_time = timestamp
            if len(self.timestamps) > 5000:
                data_slice = slice(-5000, len(self.timestamps))
            else:
                data_slice = slice(0, len(self.timestamps))
            self.data_curve.setData(self.timestamps[data_slice], self.data[data_slice])
            self.moving_avg_curve.setData(self.timestamps[data_slice], self.moving_averages[data_slice])
            # self.local_voltage_curve.setData(self.timestamps[data_slice], self.local_voltages[data_slice])
            # self.local_max_curve.setData(self.timestamps[self.local_max], self.data[self.local_max])
            # self.local_min_curve.setData(self.timestamps[self.local_min], self.data[self.local_min])
            self.signal_plot.setXRange(self.timestamps[-1] - 10, self.timestamps[-1])
            if len(self.timestamps) > 10:
                self.frequency_label.setText(f"Frequency: {int(10 / (self.timestamps[-1] - self.timestamps[-10]))} Hz")

    def update_local_attrs(self, window):
        if len(self.timestamps) == len(self.local_voltages):
            return
        
        assert len(self.timestamps) > len(self.local_voltages)

        if len(self.timestamps) == 1:
            self.local_voltages.append(self.data[0])
            self.local_max.append(0)
            self.local_min.append(0)
            self.moving_averages.append(self.data[0])
            return
        
        for last in range(len(self.local_voltages), len(self.timestamps)):            
            last_timestamp = self.timestamps[last]
            find_max = True
            if last_timestamp - self.timestamps[self.local_max[last - 1]] <= window:
                find_max = False
                if self.data[last] >= self.data[self.local_max[last - 1]]:
                    self.local_max.append(last)
                else:
                    self.local_max.append(self.local_max[last - 1])

            find_min = True
            if last_timestamp - self.timestamps[self.local_min[last - 1]] <= window:
                find_min = False
                if self.data[last] <= self.data[self.local_min[last - 1]]:
                    self.local_min.append(last)
                else:
                    self.local_min.append(self.local_min[last - 1])
            
            local_max = local_min = i = len(self.timestamps) - 1
            while i >= 0 and self.timestamps[last] - self.timestamps[i] <= window:
                if find_max and self.data[i] > self.data[local_max]:
                    local_max = i
                elif find_min and self.data[i] < self.data[local_min]:
                    local_min = i
                i -= 1

            if find_max:
                self.local_max.append(local_max)

            if find_min:
                self.local_min.append(local_min)

            self.local_voltages.append(self.data[self.local_max[-1]] - self.data[self.local_min[-1]])
            # moving_avg = (self.data[self.local_max[-1]] + self.data[self.local_min[-1]]) / 2
            # self.moving_averages.append(moving_avg)
            self.moving_averages.append(self.moving_averages[-1] * 0.998 + self.data[last] * 0.002)

if __name__ == '__main__':
    app = BMD101App()
    app.run()